# KeepAlive.py
#
# Copyright (c) 2006 Michael Hobbs
#
# This file is part of iTorrent.
#
# iTorrent is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# iTorrent is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with iTorrent; if not, write to the Free Software
# Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA  02110-1301  USA


PICK_AHEAD = 5  # Number of pieces to request ahead of the current piece.

# iTunes requires a minimum rate of about 550 bytes/sec in order to keep a
# connection open indefinitely.
ITUNES_MIN_BPS = 600  # Pad it up to 600 to allow for a little slop.
SPACES = ' ' * ITUNES_MIN_BPS


class KeepAlive(object):
	def __init__(self, writer, largestInfo, display):
		super(KeepAlive, self).__init__()
		self.writer = writer
		self.largestInfo = largestInfo
		self.currentPiece = largestInfo.begin[0]
		self.currentPieceLen = -1
		self.offset = largestInfo.begin[1]
		self.display = display

	def finish(self, torrent, callback, *args, **kwargs):
		storageWrapper = torrent._storagewrapper
		self.writer.onFlush(self.flushOut, storageWrapper, callback, args, kwargs)
		remaining = self.currentPieceLen - self.offset
		# TODO: what if we run past the end of largest file while still doing keepAlive?
		flushed, bytesSent = self.sendData(storageWrapper, remaining)
		self.currentPiece += 1
		if flushed:
			self.flushOut(storageWrapper, callback, args, kwargs)
		# end if

	def flushOut(self, storageWrapper, callback, args, kwargs):
		endPiece, endOffset = self.largestInfo.end
		flushed = True
		while flushed and self.currentPiece <= endPiece:
			self.currentPieceLen = storageWrapper._piecelen(self.currentPiece)
			if self.currentPiece == endPiece:
				self.currentPieceLen = endOffset
			self.offset = 0
			flushed, bytesSent = self.sendData(storageWrapper, self.currentPieceLen)
			self.currentPiece += 1
		if flushed and self.currentPiece > endPiece:
			self.writer.onFlush(None)
			callback(*args, **kwargs)
		# end if

	def keepAlive(self, torrent):
		# TODO: what if we run past the end of largest file while still doing keepAlive?
		storageWrapper = torrent._storagewrapper
		sendSize = ITUNES_MIN_BPS
		if self.currentPieceLen < 0:
			self.currentPieceLen = storageWrapper._piecelen(self.currentPiece)
			# With the first piece, we send out 2K of data right away to prevent
			# iTunes from spiking the CPU while waiting for data.
			sendSize = 2048
		flushed, bytesSent = self.sendData(storageWrapper, sendSize)
		if bytesSent < 0:
			return
		# Check if we need to start sending next piece to supply ITUNES_MIN_BPS
		remaining = sendSize - bytesSent
		if remaining:
			self.currentPiece += 1
			self.currentPieceLen = storageWrapper._piecelen(self.currentPiece)
			self.offset = 0
			self.sendData(storageWrapper, remaining)
		# end if

	def sendData(self, storageWrapper, numBytes):
		numBytes = min(self.currentPieceLen - self.offset, numBytes)
		if numBytes == 0:
			return (True, numBytes)
		data = storageWrapper.get_piece(self.currentPiece, self.offset, numBytes)
		if not data:
			self.display.warning('KeepAlive: No data found for piece %d' % self.currentPiece)
			return (True, -1)
		flushed = self.writer.write(data)
		self.offset += numBytes
		return (flushed, numBytes)

	def pickPiece(self, have):
		for i in range(PICK_AHEAD + 1):
			piece = self.currentPiece + i
			if piece > self.largestInfo.end[0]:
				return None
			if not have[piece]:
				return piece
			# end if
		return None


class FeedKeepAlive(object):
	"""Try to keep the connection to iTunes open as long as possible while
		prefetching the first piece(s). iTunes will eventually close the connection
		if we send it too much data, but there's not much else that we can do."""
	def __init__(self, writer, rawServer):
		super(FeedKeepAlive, self).__init__()
		self.writer = writer
		self.rawServer = rawServer
		self.stopped = False

	def start(self):
		if self.stopped:
			return
		self.rawServer.add_task(self.start, 1)
		self.writer.write(SPACES)

	def stop(self):
		self.stopped = True
